package com.jh.cloud.servlet.gateway.loadbalance;

import com.netflix.loadbalancer.BaseLoadBalancer;
import com.netflix.loadbalancer.RandomRule;
import com.netflix.loadbalancer.Server;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.scheduling.annotation.Scheduled;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class ZookeeperLoadBalancer extends BaseLoadBalancer {

    private final DiscoveryClient discoveryClient;

    private Map<String, BaseLoadBalancer> loadBalancerMap = new ConcurrentHashMap<>();

    public ZookeeperLoadBalancer(DiscoveryClient discoveryClient) {
        this.discoveryClient = discoveryClient;
        updateService();
    }

    @Override
    public Server chooseServer(Object key) {
        if (key instanceof String) {
            String serviceName = String.valueOf(key);
            BaseLoadBalancer loadBalancer = loadBalancerMap.get(serviceName);
            return loadBalancer.chooseServer(serviceName);
        }

        return super.chooseServer(key);
    }

    @Scheduled(fixedRate = 1000 * 5)
    public void updateService() {
        discoveryClient.getServices().forEach(serviceName -> {
            // 设置Rule，如果不设置的话chooseServer方法里面拿到的loadBalancer就是网关自己的loadBalancer，不是对应服务的loadBalancer
            RandomRule baseRule = new RandomRule();
            BaseLoadBalancer baseLoadBalancer = new BaseLoadBalancer();
            List<ServiceInstance> serviceInstances = discoveryClient.getInstances(serviceName);
            serviceInstances.forEach(serviceInstance -> {
                Server server = new Server(serviceInstance.isSecure() ? "https://" : "http://",
                        serviceInstance.getHost(),
                        serviceInstance.getPort());
                baseLoadBalancer.addServer(server);
            });
            // service loadBalancer link base custom rule
            baseRule.setLoadBalancer(baseLoadBalancer);
            baseLoadBalancer.setRule(baseRule);
            loadBalancerMap.put(serviceName, baseLoadBalancer);
        });
    }

}
